package LibDL.recommender;

import LibDL.core.nn.Module;
import LibDL.core.optim.Optimizer;
import LibDL.recommender.data.utils.NumUtils;
import net.librec.common.LibrecException;
import net.librec.data.structure.AbstractBaseDataEntry;
import net.librec.data.structure.BaseDataList;
import net.librec.data.structure.BaseRankingDataEntry;
import net.librec.data.structure.LibrecDataList;
import net.librec.math.structure.DataSet;
import net.librec.math.structure.SequentialAccessSparseMatrix;
import net.librec.recommender.AbstractRecommender;
import net.librec.recommender.item.KeyValue;
import net.librec.recommender.item.RecommendedList;
import org.jetbrains.annotations.Nullable;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

public abstract class LibDLRecommender extends AbstractRecommender {

    protected Module module;

    protected Optimizer optimizer;

    protected int epoch;

    protected int batch_size;

    protected int eval_per_epoch;

    protected float learning_rate;

    protected float l2;

    public final static int USE_CPU = 0;

    public final static int USE_CUDA = 1;

    protected int device;

    protected int num_items;

    protected int num_users;

    @Override
    protected void setup() throws LibrecException {
        super.setup();
        epoch = conf.getInt("libdl.epoch", 10);
        batch_size = conf.getInt("libdl.batch.size", 512);
        eval_per_epoch = conf.getInt("libdl.eval.epoch", 10);
        learning_rate = conf.getFloat("libdl.learning.rate", 0.001f);
        l2 = conf.getFloat("libdl.l2", 1e-6f);
        switch (conf.get("libdl.device", "").toLowerCase()) {
            case "cuda":
                device = 1;
                break;
            default:
                device = 0;
        }
    }

    public abstract List<Float> predict(int user_id, @Nullable List<Integer> item_ids);

    @Override
    public RecommendedList recommendRank() throws LibrecException {
        LibrecDataList<AbstractBaseDataEntry> librecDataList = new BaseDataList<>();
        for (int user_id = 0; user_id < num_users; user_id++) {
            BaseRankingDataEntry baseRankingDataEntry = new BaseRankingDataEntry(user_id);
            librecDataList.addDataEntry(baseRankingDataEntry);
        }
        return recommendRank(librecDataList);
    }

    @Override
    public RecommendedList recommendRank(LibrecDataList<AbstractBaseDataEntry> dataList) throws LibrecException {
        RecommendedList rl = new RecommendedList(dataList.size());
        ArrayList<KeyValue<Integer, Double>> row;
        int topn = this.topN;
        for (int i = 0; i < dataList.size(); i++) {
            int user_id = ((BaseRankingDataEntry) dataList.getDataEntry(i)).getUserId();
            row = new ArrayList<>(topn);
            final List<Float> pred = this.predict(user_id, null);
            int[] train = ((SequentialAccessSparseMatrix) this.getDataModel().getTrainDataSet()).row(user_id).getIndices();
            List<Integer> indices = NumUtils.argsort(pred, (o1, o2) -> pred.get(o2).compareTo(pred.get(o1)));
            indices.removeAll(IntStream.of(train).boxed().collect(Collectors.toList()));
            for (int j = 0; j < topn; j++) {
                row.add(new KeyValue<>(indices.get(j), (double) pred.get(indices.get(j))));
            }
            rl.addList(row);
        }
        return rl;
    }

    @Override
    public RecommendedList recommendRating(DataSet predictDataSet) throws LibrecException {
        return null;
    }

    @Override
    public RecommendedList recommendRating(LibrecDataList<AbstractBaseDataEntry> dataList) throws LibrecException {
        return null;
    }
}
